#!/usr/bin/env python
# $Id: netflowd.py 7557b5d4d82b 2012/12/02 09:07:16 pgurumur $
# Copyright (c) 2012 Prabhu Gurumurthy <pgurumur@gmail.com>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 7557b5d4d82b $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-12-02 09:07:16Z $"
# udpqotd - UDP message server

import datetime
import optparse
import os
import SocketServer
import struct
import sys

from venya.lib import exception, cmdparser, ip, stdlib

PORTNO = 9995

GetTime = lambda value: datetime.date.fromtimestamp(value).strftime(
      "%A %b %d %I:%M:%S %p %Y")

class FlowRecord(object):
   def __init__(self, Address):
      self._address = ip.ip(Address)
      self._netflow = None
      self._templateFlowset = None
      self._optionsTemplateFlowset = None
      self._dataFlowset = None

   def __str__(self):
      return str(self._address)

   def AddNetFlowHeader(self, Buffer):
      self._netflow = NetFlow(Buffer)

   def AddTemplateFlowSet(self, Buffer):
      if self._netflow:
         self._templateFlowset = TemplateFlowset(Buffer)

   def AddOptionsTemplateFlowSet(self, Buffer):
      self._optionsTemplateFlowset = OptionsTemplateFlowset(Buffer)

   def AddDataFlowSet(self, Buffer):
      if self._netflow and self._templateFlowset:
         self._dataFlowset = DataFlowset(Buffer, self.template)

   def parse(self):
      if self.netflow:
         self.netflow.parse()

      if self.template:
         self.template.parse()

      if self.optstemplate:
         self.optstemplate.parse()

      if self.dataset:
         self.dataset.parse()

   netflow = property(fget = lambda self: self._netflow)
   template = property(fget = lambda self: self._templateFlowset)
   optstemplate = property(fget = lambda self: self._optionsTemplateFlowset)
   dataset = property(fget = lambda self: self._dataFlowset)

class FlowHeader(object):
   def __init__(self, Buffer):
      super(FlowHeader, self).__init__()
      self._pktlen = len(Buffer)
      self._packet = Buffer

      if self._pktlen <= 20:
         self.error("invalid netflow header size, length: %d", self._pktlen)

   def __len__(self):
      return self._pktlen

   def parse(self):
      raise NotImplementedError()

   def error(self, fmt, *args):
      mesg = stdlib.sprintf(fmt, *args)
      raise ValueError(mesg)

class NetFlow(FlowHeader):
   def __init__(self, Buffer):
      super(NetFlow, self).__init__(Buffer)
      self._ver = 0
      self._count = 0
      self._uptime = 0
      self._unixtm = 0
      self._seq = 0
      self._srcid = 0
      try:
         self._header = struct.unpack(">HHIIII", self._packet[:20])
      except struct.error as err:
         self.error(str(err))
      else:
         self._ver = temp[0]
         self._count = temp[1]
         self._uptime = temp[2]
         self._unixtm = temp[3]
         self._seq = temp[4]
         self._srcid = temp[5]
         self._flowtype = struct.unpack("!H", Buffer[20:22])

   def __int__(self):
      return int(self._count)

   version = property(fget = lambda self: self._ver)
   sequence = property(fget = lambda self: self._seq)
   sourceid = property(fget = lambda self: self._srcid)
   count = property(fget = lambda self: self._count)
   flowtype = property(fget = lambda self: self._flowtype)

class TemplateFlowset(FlowHeader):
   def __init__(self, Buffer):
      super(TemplateFlowset, self).__init__(Buffer)
      self._values = {}
      self._dict = {
            1: 'IN_BYTES',
            2: 'IN_PKTS',
            3: 'FLOWS',
            4: 'PROTOCOL',
            5: 'SRC_TOS',
            6: 'TCP_FLAGS',
            7: 'L4_SRC_PORT',
            8: 'IPV4_SRC_ADDR',
            9: 'SRC_MASK',
            10: 'INPUT_SNMP',
            11: 'L4_DST_PORT',
            12: 'IPV4_DST_ADDR',
            13: 'DST_MASK',
            14: 'OUTPUT_SNMP',
            15: 'IPV4_NEXT_HOP',
            16: 'SRC_AS',
            17: 'DST_AS',
            18: 'BGP_IPV4_NEXT_HOP',
            19: 'MUL_DST_PKTS',
            20: 'MUL_DST_BYTES',
            21: 'LAST_SWITCHED',
            22: 'FIRST_SWITCHED',
            23: 'OUT_BYTES',
            24: 'OUT_PKTS',
            25: 'MIN_PKT_LENGTH',
            26: 'MAX_PKT_LENGTH',
            27: 'IPV6_SRC_ADDR',
            28: 'IPV6_DST_ADDR',
            29: 'IPV6_SRC_MASK',
            27: 'IPV6_DST_MASK',
            31: 'IPV6_FLOW_LABEL',
            32: 'ICMP_TYPE',
            33: 'MUL_IGMP_TYPE',
            34: 'SAMPLING_INTERVAL',
            35: 'SAMPLING_ALGORITHM',
            36: 'FLOW_ACTIVE_TIMEOUT',
            37: 'FLOW_INACTIVE_TIMEOUT',
            38: 'ENGINE_TYPE',
            39: 'ENGINE_ID',
            40: 'TOTAL_BYTES_EXP',
            41: 'TOTAL_PKTS_EXP',
            42: 'TOTAL_FLOWS_EXP',
            44: 'IPV4_SRC_PREFIX',
            45: 'IPV4_DST_PREFIX',
            46: 'MPLS_TOP_LABEL_TYPE',
            47: 'MPLS_TOP_LABEL_IPADDR',
            48: 'FLOW_SAMPLER_ID',
            49: 'FLOW_SAMPLER_MODE',
            50: 'FLOW_SAMPLER_RANDOM_INTERVAL',
            52: 'MIN_TTL',
            53: 'MAX_TTL',
            54: 'IPV4_IDENT',
            55: 'DST_TOS',
            56: 'IN_SRC_MAC',
            57: 'OUT_DST_MAC',
            58: 'SRC_VLAN',
            59: 'DST_VLAN',
            60: 'IP_VERSION',
            61: 'DIRECTION',
            62: 'IPV6_NEXT_HOP',
            63: 'BGP_IPV6_NEXT_HOP',
            64: 'IPV6_OPTION_HEADERS',
            70: 'MPLS_LABEL1',
            71: 'MPLS_LABEL2',
            72: 'MPLS_LABEL3',
            73: 'MPLS_LABEL4',
            74: 'MPLS_LABEL5',
            75: 'MPLS_LABEL6',
            76: 'MPLS_LABEL7',
            77: 'MPLS_LABEL8',
            78: 'MPLS_LABEL9',
            79: 'MPLS_LABEL10',
            80: 'IN_DST_MAC',
            81: 'OUT_SRC_MAC',
            82: 'IF_NAME',
            83: 'IF_DESC',
            84: 'SAMPLER_NAME',
            85: 'IN_PERM_BYTES',
            86: 'IN_PERM_PKTS',
            88: 'FRAGMENT_OFFSET',
            89: 'FORWARDING_STATUS',
            90: 'MPLS_PAL_RD',
            91: 'MPLS_PREFIX_LEN',
            92: 'SRC_TRAFFIC_INDEX',
            93: 'DST_TRAFFIC_INDEX',
            94: 'APP_DESCR',
            95: 'APP_TAG',
            96: 'APP_NAME',
            98: 'DIFFSERV_PT',
            99: 'REPLICATION_FACTOR',
            102: 'L2_PKT_SECTION_OFFSET',
            103: 'L2_PKT_SECTION_SIZE',
            104: 'L2_PKT_SECTION_DATA'
      }

   def parse(self):
      mod = self._pktlen % 2
      if mod != 0:
         self._pktlen -= 1

      count = (self._pktlen - 28) / 2
      fformat = "!%dH" %count
      count = 0
      records = struct.unpack(fformat, self._buffer[28:self._pktlen])
      for value in records:
         value = int(value)
         if (count % 2 == 0) and (value in self._dict):
            self._values[self._dict[value]] = value

         count += 1

   tvalues = property(fget = lambda self: self._values)

class DataFlowset(object):
   def __init__(self, Buffer, Template):
      super(DataFlowset, self).__init__()
      self._buffer = Buffer
      self._template = Template

class handler(SocketServer.DatagramRequestHandler):
   FlowClient = {}
   def setup(self):
      address = self.client_address[0]
      if address not in self.FlowClient:
         self.FlowClient[address] = FlowRecord(address)

   def handle(self):
      record = None
      address = self.client_address[0]
      if address in self.FlowClient:
         record = self.FlowClient[address]

      if record:
         packet = ""
         for ix in self.rfile.readlines():
            packet += ix

         pktlen = len(packet)
         if pktlen == 0:
            print("ignoring netflow header, because the size is zero!")
         elif pktlen < 20:
            version = struct.unpack("!H", packet[:2])
            print("invalid netflow header size, length: %d, version: %d" %(
               pktlen, version[0]))
         elif pktlen >= 20:
            pkthdr = NetFlow(packet)
            if not record.netflow:
               record.AddNetFlowHeader(pkthdr)

            if record.netflow.flowtype == 0:
               if record.template is None:
                  record.AddTemplateFlowSet(packet)
            elif record.netflow.flowtype == 1:
               if record.optstemplate is None:
                  record.AddOptionsTemplateFlowSet(packet)
            elif record.netflow.flowtype > 255:
               if record.dataset is None
                  record.AddDataFlowSet(packet)

            record.parse()

            print("template flowset, length: %d, template id: %d, count: %d" %(
               flowtype[1], flowtype[2], flowtype[3]))
            mod = pktlen % 2
            if mod != 0:
               pktlen -= 1

            count = (pktlen - 28) / 2
            hformat = "!%dH" %count
            records = struct.unpack(hformat, packet[28:pktlen])
            count = 0
            for value in records:
               value = int(value)
               if (count % 2 == 0) and (value in templateDict):
                  tempdict[templateDict[value]] = None

               count += 1

            print(count)
         elif flowtype[0] == 1:
            print("options template flowset, length: %d, template id: %d" %(
               flowtype[1], flowtype[2]))
            optlength = struct.unpack("!H", packet[28:30])
            print("option scope length: %d, option length: %d" %(
               flowtype[3], optlength[0]))
            records = struct.unpack("!HH", packet[30:34])
            print(records)
         elif flowtype[0] > 255:
            print("data flowset, id: %d, length: %d" %(
               flowtype[0], flowtype[1]))

            mod = pktlen % 4
            if mod != 0:
               pktlen -= mod

            count = (pktlen - 28) / 4
            dformat = "!%dI" %count
            records = struct.unpack(dformat, packet[28:pktlen])
            print(records)

      # print int(NetFlowPacketHeader(newmsg[:20]))
      # FlowHeader(newmsg[20:28], newmsg[28:])

      """
      if len(newmsg) > 20:
         strlen = len(newmsg) - 40
         # format = ">hhIIII%ds" %strlen
         format = ">hhIIIIIIIII%ds" %strlen
         temp = struct.unpack(format, newmsg)
        # print "Client %s said ``%s''" % (self.client_address[0], newmsg)
        # self.wfile.write(self.server.oldmsg)
      self.server.oldmsg = newmsg
      """


"""
s = SocketServer.UDPServer(('',PORTNO), handler)
print "Awaiting UDP messages on port %d" % PORTNO
s.oldmsg = "This is the starting message."
s.serve_forever()
"""

class Server(object):
   def __init__(self, Address = '', Port = 0):
      self._addr = None
      self._port = 0

      if Address:
         self._addr = Address

      if Port:
         self._port = int(Port)

      self._server = SocketServer.UDPServer(('', self._port), handler)

   def listen(self):
      self._server.serve_forever()

class netflowd(cmdparser.GetOpt):
   option_list = (
         optparse.make_option("-p", "--port", type = "int",
            help = "UDP port to listen, default: 9995", default = 9995,
            dest = "port"),
   ) + cmdparser.GetOpt.option_list

   def __init__(self, **Keywords):
      super(netflowd, self).__init__(**Keywords)
      self._server = None

   def handle(self):
      (options, args) = self.parse()
      if options.port:
         self._server = Server('', options.port)
         self._server.listen()

if __name__ == "__main__":
   flowd = None
   try:
      flowd = netflowd(version = 1.0, args = sys.argv)
      if len(flowd) <= 1:
         flowd.printhelp()
      else:
         flowd.handle()
   except exception.GetOptError as err:
      print(err)
      sys.exit(1)

